import os
import pandas as pd
from matplotlib.figure import Figure
from typing import Tuple, Dict, List, Union
from src.core.academic_atmosphere.generate_academic_atmosphere_graph import (
    create_score_plot,
    create_attendance_plot,
    create_consumption_scatter,
    create_campus_boxplot
)
from src.common.descriptors.load_data import (
    load_stu_data, 
    load_grade_data, 
    load_kaoqin_data, 
    load_kaoqin_type_data, 
    load_consumption_data
)

def load_and_preprocess() -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    df_student = load_stu_data()
    df_kaoqin = load_kaoqin_data()
    df_kaoqin_type = load_kaoqin_type_data()
    df_grade = load_grade_data()
    df_consump = load_consumption_data()

    df_kaoqin = pd.merge(df_kaoqin, df_kaoqin_type, 
                        left_on='control_task_order_id',
                        right_on='control_task_order_id',
                        how='left')

    valid_scores = df_grade[df_grade['mes_Score'] > 0]

    df_consump['DealDate'] = pd.to_datetime(df_consump['DealTime']).dt.date
    daily_consump = df_consump.groupby(['bf_StudentID', 'DealDate'])['MonDeal'].sum().reset_index()
    
    return df_student, df_kaoqin, valid_scores, daily_consump

def analyze_class(df_student: pd.DataFrame, df_kaoqin: pd.DataFrame, valid_scores: pd.DataFrame, daily_consump: pd.DataFrame, cla_name: str) -> Dict[str, Union[str, int, float, Dict[str, float], Dict[str, int]]]:
    class_students = df_student[df_student['cla_Name'] == cla_name]
    stu_ids = class_students['bf_StudentID'].unique()
    
    class_kaoqin = df_kaoqin[df_kaoqin['cla_Name'] == cla_name]
    kaoqin_stats = class_kaoqin.groupby('control_task_name').size().sort_values(ascending=False)
    
    class_grades = valid_scores[valid_scores['mes_StudentID'].isin(stu_ids)]
    avg_score = class_grades['mes_Score'].mean()
    score_dist = class_grades.groupby('mes_sub_name')['mes_Score'].mean()
    
    class_consump = daily_consump[daily_consump['bf_StudentID'].isin(stu_ids)]
    avg_daily = class_consump.groupby('bf_StudentID')['MonDeal'].apply(lambda x: x.abs().mean()).mean()
    
    return {
        "班级名称": cla_name,
        "学生人数": len(class_students),
        "平均成绩": round(avg_score, 1) if not pd.isna(avg_score) else 0,
        "主要违纪类型": kaoqin_stats.index[0] if not kaoqin_stats.empty else "无",
        "日均消费(元)": round(abs(avg_daily), 1),
        "成绩分布": score_dist.to_dict(),
        "考勤统计": kaoqin_stats.to_dict()
    }

def generate_academic_atmosphere_data(df_report: pd.DataFrame, analysis_results: List[Dict[str, Union[str, int, float, Dict[str, float], Dict[str, int]]]]):
    os.makedirs('data', exist_ok=True)

    base_df = df_report[['班级名称', '学生人数', '平均成绩', '日均消费(元)', '主要违纪类型']]
    base_df = base_df.rename(columns={'班级名称': '班级'})

    subject_data = []
    seen_subjects = set()
    unique_subjects = []
    for record in analysis_results:
        class_name = record['班级名称']
        for subj, score in record['成绩分布'].items():
            subject_data.append({
                '班级': class_name,
                '科目': subj,
                '平均分': round(score, 1)
            })
            if subj not in seen_subjects:
                seen_subjects.add(subj)
                unique_subjects.append(subj)

    attend_data = []
    seen_attend = set()
    unique_attend = []
    for record in analysis_results:
        class_name = record['班级名称']
        for kq_type, count in record['考勤统计'].items():
            attend_data.append({
                '班级': class_name,
                '考勤类型': kq_type,
                '次数': count
            })
            if kq_type not in seen_attend:
                seen_attend.add(kq_type)
                unique_attend.append(kq_type)

    final_df = base_df.copy()

    if subject_data:
        subject_pivot = pd.DataFrame(subject_data).pivot(
            index='班级', 
            columns='科目',
            values='平均分'
        ).reset_index()
        final_df = final_df.merge(subject_pivot, on='班级', how='left')

    if attend_data:
        attend_pivot = pd.DataFrame(attend_data).pivot(
            index='班级',
            columns='考勤类型',
            values='次数'
        ).reset_index()
        final_df = final_df.merge(attend_pivot, on='班级', how='left')

    base_cols = ['班级', '学生人数', '平均成绩']
    middle_cols = ['日均消费(元)', '主要违纪类型']

    subject_cols = unique_subjects if subject_data else []
    attend_cols = unique_attend if attend_data else []

    ordered_cols = base_cols + subject_cols + middle_cols + attend_cols

    final_df = final_df[ordered_cols].fillna(0)
    final_df = final_df.sort_values('班级').reset_index(drop=True)

    for col in attend_cols:
        final_df[col] = final_df[col].astype(int)

    final_df.to_csv('data/academic_atmosphere.csv', index=False, encoding='utf_8_sig')

def analyze_style() -> Tuple[pd.DataFrame, Figure, Figure, Figure, Figure]:
    df_student, df_kaoqin, valid_scores, daily_consump = load_and_preprocess()

    all_classes = df_student['cla_Name'].dropna().unique()
    analysis_results = []
    for cla in all_classes:
        if str(cla).startswith(('白', '东')):
            result = analyze_class(df_student, df_kaoqin, valid_scores, daily_consump, cla)
            analysis_results.append(result)

    df_report = pd.DataFrame(analysis_results)

    try:
        academic_atmosphere = pd.read_csv('data/academic_atmosphere.csv')
    except:
        generate_academic_atmosphere_data(df_report, analysis_results)
        academic_atmosphere = pd.read_csv('data/academic_atmosphere.csv')
    
    fig1 = create_score_plot(df_report)
    fig2 = create_attendance_plot(analysis_results)
    fig3 = create_consumption_scatter(df_report)
    fig4 = create_campus_boxplot(df_report)
    
    return academic_atmosphere, fig1, fig2, fig3, fig4